# -*- coding: utf-8 -*-
# @Time    : 2020/5/12 下午6:27
# @Author  : tsl
# @Email   : mymailwith163@163.com
# @File    : mk_train_val_data.py
# @Software: PyCharm

import os,cv2
from videoLabel import videoLabel as vl

def remove_small_labels(src,dst):
    for file in os.listdir(src):
        if(file.endswith(".txt")):
            if(not file.startswith("class")):
                print(file)
                
                img = cv2.imread(os.path.join(src,file.replace(".txt",".jpg")))
                _,iw,ih = img.shape[::-1]
                dstLabel = open(os.path.join(dst,file),"w")
                with open(os.path.join(src,file),"r") as f:
                    lines = f.readlines()
                    for line in lines:
                        label,rx,ry,rw,rh = map(float,line.rstrip().split(" "))
                        x,y,w,h = vl.convertR(rx,ry,rw,rh,iw,ih)
                        if(w>10 and h>10):
                            rx_,ry_,rw_,rh_ = vl.convert([iw,ih],[x,x+w,y,y+h])
                            dstLabel.write(" ".join(map(str,[int(label),rx_,ry_,rw_,rh_]))+"\n" )
                dstLabel.close()


def mk_train_val(src,train_txt,val_txt):
    train_rat = 0.9
    train_num = int(len(os.listdir(src))/2 *train_rat)
    for num,file in enumerate(os.listdir(src)):
        if(not file.endswith(".txt")):
            if(int(num/2)<train_num):

                train_txt.write(os.path.join(src,file)+"\n")
            else:

                val_txt.write(os.path.join(src,file)+"\n")

def run_mk_train_val():
    src = ["/mnt/data/Pyproject/head/head/SCUT_HEAD_Part_B/train_data",
           "/mnt/data/Pyproject/head/train_head_1",
           "/mnt/data/Pyproject/head/train_head_2"]

    train_txt = open("train.txt","w")
    val_txt = open("val.txt","w")

    for s in src:
        mk_train_val(s,train_txt,val_txt)

    train_txt.close()
    val_txt.close()

def run_remove_small_labels():
    src = "/mnt/data/Pyproject/head/test_head/5_img"
    dst="/mnt/data/Pyproject/head/test_head/5_txt"
    remove_small_labels(src,dst)


def main():
    # run_mk_train_val()
    run_remove_small_labels()


if __name__ == "__main__":
    main()